import urllib.request
import shutil
import gzip
import os

test_label_url = 'http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz'
test_image_url = 'http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz'
train_label_url = 'http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz'
train_image_url = 'http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz'

def unzip_and_copy_to(in_file_path: str, out_file_path: str) -> None:
    with gzip.open(in_file_path, 'rb') as f_in:
        with open(out_file_path, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)

shutil.rmtree('./test-set', ignore_errors=True)
os.mkdir('./test-set')

print('Downloading test set labels...')
filepath, _ = urllib.request.urlretrieve(test_label_url)
unzip_and_copy_to(filepath, './test-set/labels')

print('Downloading test set images...')
filepath, _ = urllib.request.urlretrieve(test_image_url)
unzip_and_copy_to(filepath, './test-set/images')

shutil.rmtree('./train-set', ignore_errors=True)
os.mkdir('./train-set')

print('Downloading train set labels...')
filepath, _ = urllib.request.urlretrieve(train_label_url)
unzip_and_copy_to(filepath, './train-set/labels')

print('Downloading train set images...')
filepath, _ = urllib.request.urlretrieve(train_image_url)
unzip_and_copy_to(filepath, './train-set/images')
